# imports
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.datasets as datasets
import torchvision.transforms as transforms

# create fully connected networks
class FCN(nn.Module):
    def __init__(self, input_size, num_classes):
        super(FCN, self).__init__()
        self.fc1 = nn.Linear(input_size, 50)
        self.fc2 = nn.Linear(50, num_classes)

    def forward(self, x):
        return self.fc2(F.relu(self.fc1(x)))


# train
def train(model, loss_fn, optimizer, dataloader, epochs):
    model.train()
    print("=" * 10 + "training" + "=" * 10)
    for epoch in range(epochs):
        correct_samples = 0
        total_samples = 0
        total_loss = 0

        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(device), target.to(device)
            data = data.reshape(data.shape[0], -1)

            result = model(data)
            loss = loss_fn(result, target)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()
            total_samples += target.shape[0]
            correct_samples += target.eq(result.argmax(-1)).sum().item()
        print(f"epoch: {epoch}. loss: {total_loss/ total_samples}, accuracy: {correct_samples / total_samples * 100:.2f}% ")


def test(model, loss_fn, dataloader):
    model.eval()
    print("=" * 10 + "testing" + "=" * 10)
    with torch.no_grad():
        total_samples = 0
        correct_samples = 0
        total_loss = 0
        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(device), target.to(device)
            data = data.reshape(data.shape[0], -1)
            result = model(data)
            loss = loss_fn(result, target)

            total_loss += loss.item()
            total_samples += target.shape[0]
            correct_samples += target.eq(result.argmax(-1)).sum().item()
        print(f"loss: {loss / total_samples}, accuracy: {correct_samples / total_samples * 100:.2f}% ")

# setup device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# hyperparameters
batch_size = 128
input_size = 784
num_classes = 10
learning_rate = 0.001
epochs = 10

# load data
train_data = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)
test_data = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())
test_loader = torch.utils.data.DataLoader(dataset=test_data, batch_size=batch_size, shuffle=True)

# initialize network, loss, optimizer
model = FCN(input_size, num_classes).to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

if __name__ == '__main__':
    train(model, loss_fn, optimizer, train_loader, epochs)
    test(model, loss_fn, test_loader)





